import numpy as np
from pprint import pprint
import matplotlib.pyplot as plt


def p2_1():
    # plotting the data
    data1 = np.loadtxt("ex1data1.txt", delimiter=",")
    data1 = data1[np.argsort(data1[:, 0])]
    x = data1[:, 0]
    y = data1[:, 1]
    plt.xlabel("Profit in $10,000s")
    plt.ylabel("Population of City in 10,000s")
    plt.plot(x, y, "k")
    plt.show()


def J_theta(theta, x, y):
    m = len(y)

    return np.sum((theta @ x - y) ** 2) / (2 * m)


def updateTheta(theta, x, y, lr):
    m = len(y)
    x0 = x[0, :]
    x1 = x[1, :]

    tmp = (theta @ x - y)
    return np.array([
        theta[0] - lr * (1 / m) * np.sum(tmp * x0),
        theta[1] - lr * (1 / m) * np.sum(tmp * x1)
    ])


def linearRegression():
    data1 = np.loadtxt("ex1data1.txt", delimiter=",")
    x = data1[:, 0]
    x = np.vstack((np.ones(len(x)), x))
    y = data1[:, 1]

    theta = np.zeros(2)
    lr = 0.023
    iterations = 1500

    with open("theta.csv", "w+") as theta_f, \
            open("cost.csv", "w+") as cost_f:
        for i in range(iterations):
            theta = updateTheta(theta, x, y, lr)
            cost = J_theta(theta, x, y)
            theta_f.write("{},{},{}\n".format(i, theta[0], theta[1]))
            cost_f.write("{},{}\n".format(i, cost))
            print("iteration: {}, theta: [{}, {}], cost: {}".format(i, theta[0], theta[1], cost))


def drawCosts():
    costs = np.loadtxt("cost.csv", delimiter=",")
    x = costs[:, 0]
    y = costs[:, 1]

    plt.plot(x, y, "k")
    plt.show()


def drawResult():
    data1 = np.loadtxt("ex1data1.txt", delimiter=",")
    data1 = data1[np.argsort(data1[:, 0])]
    x = data1[:, 0]
    y = data1[:, 1]
    plt.xlabel("Profit in $10,000s")
    plt.ylabel("Population of City in 10,000s")
    plt.plot(x, y, "rx")

    theta = np.loadtxt("theta.csv", delimiter=",")
    theta = theta[-1, 1:]
    theta0, theta1 = theta

    _y = theta0 + theta1 * x
    plt.plot(x, _y, "k")
    plt.show()


if __name__ == '__main__':
    linearRegression()
    drawCosts()
    drawResult()
